from select import *
from socket import *


class TCPServer:
    def __init__(self, host="0.0.0.0", port=2816):
        self.host = host
        self.port = port
        self.addr = (self.host, self.port)
        self.sock = self.create_socket()

    def create_socket(self):
        sock = socket()
        sock.bind(self.addr)
        sock.listen(5)
        sock.setblocking(False)
        return sock


class EpollServer:
    def __init__(self):
        self.ep = epoll()
        self.tcp_sock = TCPServer()
        self.map = {self.tcp_sock.sock.fileno(): self.tcp_sock.sock}

    def main(self):
        self.ep.register(self.tcp_sock.sock, EPOLLIN)
        while True:
            events = self.ep.poll()
            # 遍历就绪的IO
            for fd, event in events:
                if self.map[fd] is self.tcp_sock.sock:
                    self.connect(fd)
                else:
                    self.dispose_data(fd)

    def dispose_data(self, fd):
        data = self.map[fd].recv(1024)
        if not data:
            self.ep.unregister(fd)
            self.map[fd].close()
            del self.map[fd]
            return
        print(data.decode())
        self.map[fd].send(b"ok")

    def connect(self, fd):
        conn, addr = self.map[fd].accept()
        print("Connect from", addr)
        conn.setblocking(False)
        self.ep.register(conn, EPOLLIN)
        self.map[conn.fileno()] = conn


if __name__ == '__main__':
    ep = EpollServer()
    ep.main()
